import tushare as ts
from pymongo import MongoClient
import json
import datetime

client = MongoClient('mongodb://localhost:27017')
database = client['stocks']

# 获取股票列表集合
stock_list_collection = database['stock_list']
stock_list_collection.create_index('code', unique=True)

stock_list_df = ts.get_stock_basics()

stock_list_data = []
# 遍历股票列表并且存入数据库
for index, row in stock_list_df.iterrows():  # 获取每行的index、row
    item = json.loads(row.to_json())
    item['code'] = index
    if item['code'][0] == '0' or item['code'][0] == '3':
        item['exchange'] = 'SZ'
    if item['code'][0] == '6':
        item['exchange'] = 'SH'

    ipo_date = str(item['timeToMarket'])
    if len(ipo_date) == 8:
        ipo_date = datetime.datetime.strptime(ipo_date + ' 0:0:0', '%Y%m%d %H:%M:%S')
    else:
        ipo_date = ''
    stock_list_data.append({
        'name': item['name'],
        'code': item['code'],
        'exchange': item['exchange'],
        'industry': item['industry'],
        'category': '',
        'concept': '',
        'ipo_date': ipo_date
    })

stock_list_collection.drop()
stock_list_collection.insert_many(stock_list_data)
print('stock_list 成功插入', len(stock_list_df), '条数据')
